import os
import cv2
from torch.utils.data import Dataset
from torchvision.transforms import (Compose,Resize,RandomHorizontalFlip,RandomRotation,Normalize
,ColorJitter,RandomGrayscale,RandomApply,GaussianBlur,ToTensor)
from PIL import Image

mode = "train"
file ="training"

image_size = (128, 128)
# 数据增强和预处理
train_tfm = Compose([
            Resize(size=image_size),
            # 随机水平翻转 (数据增强)
            RandomHorizontalFlip(),
            RandomRotation(15),
            # 颜色抖动 (颜色增强)
            ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
            # 随机灰度化 (颜色增强)
            RandomGrayscale(p=0.1),
            # 随机高斯模糊 (图像质量增强)
            RandomApply([GaussianBlur(kernel_size=3, sigma=(0.1, 2.0))], p=0.3),
            ToTensor(),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

test_tfm = Compose([
            Resize(size=image_size),
            ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
            ToTensor(),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])


class FoodDataset(Dataset):
    def __init__(self, path, mode='train', transform=None):
        self.mode = mode
        self.transform = transform
        self.filenames = []
        self.labels = []

        for file in os.listdir(path):
            if file.endswith('.jpg'):
                self.filenames.append(os.path.join(path, file))
                if mode != 'test':
                    self.labels.append(int(file.split('_')[0]))
                else:
                    self.labels.append(0)

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, idx):
        image = Image.open(self.filenames[idx]).convert('RGB')
        if self.transform:
            image = self.transform(image)

        if self.mode != 'test':
            return image, self.labels[idx]
        else:
            return image, os.path.basename(self.filenames[idx])

if __name__ == "__main__":
    dataset = FoodDataset(f'food11/versions/1/{file}', mode=mode, transform=train_tfm)
    print(dataset[0][0].shape)

